#include "stdafx.h"
#include "h264_encoder.h"
#ifdef ENCODE_LOG
#include "DebugUtils.h"
#endif // DEBUG


H264_encoder::H264_encoder() {
    pCodecCtx = NULL;
    pCodec = NULL;
    pFrame = NULL;
    frame_buffer = NULL;
    nWidth = 1920;
    nHeight = 1080;
    nFrameSize = nWidth * nHeight * 3 / 2;
    nSize = nWidth * nHeight;
    pPacket = NULL;
    m_nImageIndex = 0;
    m_bInit = FALSE;
    m_dbCompVal = 0.6;
}

H264_encoder::~H264_encoder() {
    avcodec_close(pCodecCtx);
    avcodec_free_context(&pCodecCtx);
    av_frame_free(&pFrame);
    av_packet_free(&pPacket);
    pCodecCtx = NULL;
    pCodec = NULL;
    pFrame = NULL;
    pPacket = NULL;

}

int H264_encoder::init(int nW, int nH) {
    nWidth = nW;
    nHeight = nH;
    nFrameSize = nWidth * nHeight * 3 / 2;
    m_nImageIndex = 0;

    pCodec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!pCodec) {
        return -1;
    }

    pCodecCtx = avcodec_alloc_context3(NULL);
    if (!pCodecCtx) {
        return -1;
    }

    pCodecCtx->codec_type = AVMEDIA_TYPE_VIDEO; // 编码器类型
    pCodecCtx->codec_id = AV_CODEC_ID_H264;
    pCodecCtx->pix_fmt = AV_PIX_FMT_YUV420P;    // 像素格式
    
    // 视频帧大小，单位为像素
    pCodecCtx->width = nWidth;
    pCodecCtx->height = nHeight;

    // 视频的帧率，用分数表示，num/den
    pCodecCtx->time_base.num = 1;
    pCodecCtx->time_base.den = FPS;

    // 码率控制
    pCodecCtx->bit_rate = BIT_RATE;
    pCodecCtx->bit_rate_tolerance = 200000;

    // 量化系数
    pCodecCtx->qmin = QMIN;
    pCodecCtx->qmax = QMAX;
    pCodecCtx->max_qdiff = 4; // 所有帧的最大Q值差距
    //pCodecCtx->qblur = 0.0;   // 随时间推移的 qscale 平滑量 
    pCodecCtx->qcompress = m_dbCompVal;

    pCodecCtx->gop_size = 12;    // 每个分组的图片张数，即两个I帧的间隔
    pCodecCtx->max_b_frames = 3; // B帧的最大连续个数
    //pCodecCtx->me_range = 12;    // 运动估计区间
    //pCodecCtx->thread_count = 4; // 可用线程数

    AVDictionary* options = NULL;
    av_dict_set(&options, "preset", "slow", 0);
    av_dict_set(&options, "tune", "zerolatency", 0);
    int ret = avcodec_open2(pCodecCtx, pCodec, &options);
    if (ret != 0)
        return -2;
    av_dict_free(&options);

    pFrame = av_frame_alloc();
    pFrame->format = pCodecCtx->pix_fmt;
    pFrame->width = pCodecCtx->width;
    pFrame->height = pCodecCtx->height;
    int frame_size = avpicture_get_size(pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height);
    frame_buffer = (uint8_t*)av_malloc(frame_size);
    if (!frame_buffer)
        return -1;
    ret = avpicture_fill((AVPicture*)pFrame, frame_buffer, pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height);
    if (ret < 0)
        return -2;

    nSize = pCodecCtx->width * pCodecCtx->height;
    pPacket = av_packet_alloc();

    m_bInit = TRUE;
    return 0;
}

BOOL H264_encoder::EncodeFun(BYTE* pSrcBuf, BYTE* pDstBuf, int& nEnSize) {
    memcpy(pFrame->data[0], pSrcBuf, nSize);
    memcpy(pFrame->data[1], pSrcBuf + nSize, nSize / 4);
    memcpy(pFrame->data[2], pSrcBuf + nSize + nSize / 4, nSize / 4);

    int got_frame = 0;
    int ret = avcodec_encode_video2(pCodecCtx, pPacket, pFrame, &got_frame);
    if (ret < 0)
        return FALSE;
    if (got_frame == 1) {
        m_nImageIndex++;
        pPacket->pts = (long)m_nImageIndex * 1000 / FPS;

        memcpy(pDstBuf, pPacket->data, pPacket->size);
        nEnSize = pPacket->size;

    #ifdef ENCODE_LOG
        char* type = NULL;
        if (pFrame->pict_type == AV_PICTURE_TYPE_I) {
            type = "I";
        } else if (pFrame->pict_type == AV_PICTURE_TYPE_B) {
            type = "B";
        } else if (pFrame->pict_type == AV_PICTURE_TYPE_P) {
            type = "P";
        } else {
            type = "?";
        }
        stopTime = clock();
        writeLog("[Encode] Frame type: %s, Frame length: %d, Time cost: %dms\n", type, nEnSize, stopTime - startTime);
        startTime = stopTime;
    #endif

        av_free_packet(pPacket); // 只是将packet.data赋值为NULL(防止指针悬空)和packet.size = 0(以便进行下一次编码)
        return TRUE;
    }
    return FALSE;
}


BOOL H264_encoder::GetInitSts() {
    return m_bInit;
}

// 设置视频压缩率
void H264_encoder::SetCompVal(float dbCompVal) {
    if (NULL != pCodecCtx) {
        m_dbCompVal = dbCompVal;
        pCodecCtx->qcompress = dbCompVal;
    }
}